package com.spring.dynamicdatasource.config;

import com.spring.dynamicdatasource.utils.Keys;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpSession;
import javax.sql.DataSource;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;

/**
 * 动态数据源
 *
 * @author HouKunLin
 * @date 19-9-9 下午11:25
 */
@Primary
@Component
public class DynamicDataSource extends AbstractRoutingDataSource {

    private Map<Object, Object> targetDataSources = new LinkedHashMap<>();

    @Autowired
    private HttpSession session;

    @Override
    protected String determineCurrentLookupKey() {
        /*
        这里把当前的数据源信息存在了当前的请求会话里面，因此需要使用SESSION对象，当然也可以根据需要存在ThreadLocal当前线程里面
        但是使用ThreadLocal当前线程需要自己去设置值
         */
        return String.valueOf(session.getAttribute(Keys.datasourcesKey));
    }

    @Override
    protected DataSource determineTargetDataSource() {
        Object o = session.getAttribute(Keys.datasources);
        if (o instanceof DataSource) {
            return (DataSource) o;
        }

        return super.determineTargetDataSource();
    }

    /**
     * 重写该方法，自己维护一个数据源的列表
     *
     * @param targetDataSources 新的数据源列表
     */
    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
        afterPropertiesSet();
    }

    public Map<Object, Object> getTargetDataSources() {
        return targetDataSources;
    }

    /**
     * 这个方法必须重写，主要是为了初始化super.setTargetDataSources这个方法的值
     */
    @Override
    public void afterPropertiesSet() {
        super.setTargetDataSources(this.targetDataSources);
        super.afterPropertiesSet();
    }

    public void put(Object key, DataSource value) {
        if (key == null) {
            return;
        }
        targetDataSources.put(key, value);
        afterPropertiesSet();
    }

    public void putAll(Map<Object, Object> map) {
        if (map == null) {
            return;
        }
        targetDataSources.putAll(map);
        afterPropertiesSet();
    }

    public void remove(Object key) {
        if (key == null) {
            return;
        }
        targetDataSources.remove(key);
        afterPropertiesSet();
    }

    public void remove(DataSource dataSource) {
        Object removeKey = null;
        Set<Map.Entry<Object, Object>> entries = targetDataSources.entrySet();
        for (Map.Entry<Object, Object> entry : entries) {
            if (entry.getValue() == dataSource) {
                removeKey = entry.getKey();
                break;
            }
        }
        remove(removeKey);
    }
}
